from typing import TypedDict, Annotated
from langgraph.graph import StateGraph
from nodes import *

class AgentState(TypedDict):
    step1_output: str
    step2_output: str
    step3_output: str
    step4_output: str
    step5_output: str
    step6_output: str
    step7_output: str
    step8_output: str
    step9_output: str

class BaseNode:
    def run(self, state: AgentState) -> AgentState:
        raise NotImplementedError

class NovelWorkflow:
    def __init__(self):
        self.graph = StateGraph(AgentState)
        self._build_graph()

    def _build_graph(self):
        # 添加所有步骤节点
        nodes = {
            'step1': Step1Node(),
            'step2': Step2Node(),
            'step3': Step3Node(),
            'step4': Step4Node(),
            'step5': Step5Node(),
            'step6': Step6Node(),
            'step7': Step7Node(),
            'step8': Step8Node(),
            'step9': Step9Node()
        }

        for name, node in nodes.items():
            self.graph.add_node(name, node.run)

        # 设置线性执行流程
        self.graph.add_edge('step1', 'step2')
        self.graph.add_edge('step2', 'step3')
        self.graph.add_edge('step3', 'step4')
        self.graph.add_edge('step4', 'step5')
        self.graph.add_edge('step5', 'step6')
        self.graph.add_edge('step6', 'step7')
        self.graph.add_edge('step7', 'step8')
        self.graph.add_edge('step8', 'step9')

        self.graph.set_entry_point('step1')
        self.graph.set_finish_point('step9')

    def run(self):
        app = self.graph.compile()
        return app.invoke({'step1_output': ''})